#include "xsan_disk_manager.h"
#include "xsan_bdev.h"       // For xsan_bdev_list_get_all, xsan_bdev_info_t
#include "xsan_memory.h"   // For XSAN_MALLOC, XSAN_FREE, XSAN_CALLOC
#include "xsan_list.h"     // Using our own list for managing disks and groups
#include "xsan_list_iter.h" // For xsan_list_iter_t
#include "xsan_string_utils.h" // For xsan_strcpy_safe
#include "xsan_log.h"      // For XSAN_LOG_INFO, XSAN_LOG_ERROR, etc.
#include "xsan_metadata_store.h" // For RocksDB wrapper
#include "json-c/json.h"   // For JSON processing
#include "xsan_error.h"

#include "spdk/uuid.h"     // For spdk_uuid_generate, spdk_uuid_compare, spdk_uuid_is_null, spdk_uuid_fmt_lower
#include "spdk/bdev.h"     // For spdk_bdev_open_ext, spdk_bdev_close
#include "spdk/thread.h"   // For spdk_get_thread
#include <pthread.h>       // For pthread_mutex_t

// --- Defines for RocksDB Keys ---
#define XSAN_DISK_META_PREFIX "d:"
#define XSAN_DISK_GROUP_META_PREFIX "g:"

// The internal structure for the disk manager
struct xsan_disk_manager {
    xsan_list_t *managed_disks;       // List of xsan_disk_t structures
    xsan_list_t *managed_disk_groups; // List of xsan_disk_group_t structures
    pthread_mutex_t lock;             // Mutex to protect access to the lists
    bool initialized;                 // Flag to indicate if the manager is initialized
    xsan_metadata_store_t *md_store;  // Handle to the metadata store (RocksDB)
    char metadata_db_path[XSAN_MAX_PATH_LEN]; // Path to the DB for this manager
};

static xsan_disk_manager_t *g_xsan_disk_manager_instance = NULL;

// Forward declarations for static helpers (commented out as not implemented yet)
// static xsan_error_t xsan_disk_manager_load_metadata(xsan_disk_manager_t *dm);
// static xsan_error_t xsan_disk_manager_save_disk_meta(xsan_disk_manager_t *dm, xsan_disk_t *disk);
// static xsan_error_t xsan_disk_manager_save_group_meta(xsan_disk_manager_t *dm, xsan_disk_group_t *group);
// static xsan_error_t xsan_disk_manager_delete_disk_meta(xsan_disk_manager_t *dm, xsan_disk_id_t disk_id);
// static xsan_error_t xsan_disk_manager_delete_group_meta(xsan_disk_manager_t *dm, xsan_group_id_t group_id);

// static xsan_error_t _xsan_disk_to_json_string(const xsan_disk_t *disk, char **json_string_out);
// static xsan_error_t _xsan_json_string_to_disk(const char *json_string, xsan_disk_t **disk_out);
// static xsan_error_t _xsan_disk_group_to_json_string(const xsan_disk_group_t *group, char **json_string_out);
// static xsan_error_t _xsan_json_string_to_disk_group(const char *json_string, xsan_disk_manager_t *dm, xsan_disk_group_t **group_out);


static void _xsan_internal_disk_destroy_cb(void *disk_data) {
    if (disk_data) {
        xsan_disk_t *disk = (xsan_disk_t *)disk_data;
        char xsan_id_str[SPDK_UUID_STRING_LEN];
        spdk_uuid_fmt_lower(xsan_id_str, sizeof(xsan_id_str), (struct spdk_uuid*)&disk->id.data[0]);
        XSAN_LOG_DEBUG("Disk Manager: Destroying xsan_disk (bdev: %s, XSAN_ID: %s)", disk->bdev_name, xsan_id_str);
        if (disk->bdev_descriptor) {
            spdk_bdev_close(disk->bdev_descriptor);
            disk->bdev_descriptor = NULL;
        }
        XSAN_FREE(disk);
    }
}

static void _xsan_internal_disk_group_destroy_cb(void *group_data) {
    if (group_data) {
        xsan_disk_group_t *group = (xsan_disk_group_t *)group_data;
        char group_id_str[SPDK_UUID_STRING_LEN];
        spdk_uuid_fmt_lower(group_id_str, sizeof(group_id_str), (struct spdk_uuid*)&group->id.data[0]);
        XSAN_LOG_DEBUG("Disk Manager: Destroying xsan_disk_group (Name: '%s', ID: %s)", group->name, group_id_str);
        XSAN_FREE(group);
    }
}

xsan_error_t xsan_disk_manager_init(xsan_disk_manager_t **dm_instance_out) {
    // Path should ideally come from a global config. Using a default for now.
    // Example: data_dir from xsan_config + "/disk_manager_db"
    const char *default_db_path_suffix = "xsan_meta_db/disk_manager";
    char actual_db_path[XSAN_MAX_PATH_LEN];

    // TODO: Integrate with a proper config system to get base data path
    // For now, using a relative path. Ensure "./xsan_meta_db" directory exists.
    // In a real setup: snprintf(actual_db_path, sizeof(actual_db_path), "%s/%s", xsan_get_data_dir(), default_db_path_suffix);
    snprintf(actual_db_path, sizeof(actual_db_path), "./%s", default_db_path_suffix);


    if (g_xsan_disk_manager_instance != NULL) {
        XSAN_LOG_WARN("XSAN Disk Manager already initialized.");
        if (dm_instance_out) *dm_instance_out = g_xsan_disk_manager_instance;
        return XSAN_OK;
    }
    XSAN_LOG_INFO("Initializing XSAN Disk Manager (DB path: %s)...", actual_db_path);
    xsan_disk_manager_t *dm = (xsan_disk_manager_t *)XSAN_MALLOC(sizeof(xsan_disk_manager_t));
    if (!dm) {
        XSAN_LOG_ERROR("Failed to allocate memory for XSAN Disk Manager.");
        if (dm_instance_out) *dm_instance_out = NULL;
        return XSAN_ERROR_OUT_OF_MEMORY;
    }
    memset(dm, 0, sizeof(xsan_disk_manager_t));
    xsan_strcpy_safe(dm->metadata_db_path, actual_db_path, XSAN_MAX_PATH_LEN);

    dm->managed_disks = xsan_list_create(_xsan_internal_disk_destroy_cb);
    dm->managed_disk_groups = xsan_list_create(_xsan_internal_disk_group_destroy_cb);

    if (!dm->managed_disks || !dm->managed_disk_groups) {
        XSAN_LOG_ERROR("Failed to create lists for Disk Manager.");
        if(dm->managed_disks) xsan_list_destroy(dm->managed_disks);
        if(dm->managed_disk_groups) xsan_list_destroy(dm->managed_disk_groups);
        XSAN_FREE(dm);
        if (dm_instance_out) *dm_instance_out = NULL;
        return XSAN_ERROR_OUT_OF_MEMORY;
    }

    if (pthread_mutex_init(&dm->lock, NULL) != 0) {
        XSAN_LOG_ERROR("Failed to initialize mutex for Disk Manager.");
        xsan_list_destroy(dm->managed_disk_groups);
        xsan_list_destroy(dm->managed_disks);
        XSAN_FREE(dm);
        if (dm_instance_out) *dm_instance_out = NULL;
        return XSAN_ERROR_SYSTEM;
    }

    dm->md_store = xsan_metadata_store_open(dm->metadata_db_path, true);
    if (!dm->md_store) {
        XSAN_LOG_ERROR("Failed to open/create metadata store at '%s' for Disk Manager.", dm->metadata_db_path);
        pthread_mutex_destroy(&dm->lock);
        xsan_list_destroy(dm->managed_disk_groups);
        xsan_list_destroy(dm->managed_disks);
        XSAN_FREE(dm);
        if (dm_instance_out) *dm_instance_out = NULL;
        return XSAN_ERROR_STORAGE_GENERIC;
    }
    XSAN_LOG_INFO("Metadata store opened for Disk Manager at '%s'.", dm->metadata_db_path);

    dm->initialized = true;
    g_xsan_disk_manager_instance = dm;
    if (dm_instance_out) *dm_instance_out = g_xsan_disk_manager_instance;

    // xsan_disk_manager_load_metadata(dm); // Load existing metadata - TODO: implement this function

    XSAN_LOG_INFO("XSAN Disk Manager initialized. Call scan_and_register_bdevs() to reconcile with live bdevs.");
    return XSAN_OK;
}

void xsan_disk_manager_fini(xsan_disk_manager_t **dm_ptr) {
    xsan_disk_manager_t *dm_to_fini = NULL;
    if (dm_ptr && *dm_ptr) dm_to_fini = *dm_ptr;
    else if (g_xsan_disk_manager_instance) dm_to_fini = g_xsan_disk_manager_instance;

    if (!dm_to_fini || !dm_to_fini->initialized) {
        if (dm_ptr) *dm_ptr = NULL;
        if (dm_to_fini == g_xsan_disk_manager_instance) g_xsan_disk_manager_instance = NULL;
        return;
    }

    XSAN_LOG_INFO("Finalizing XSAN Disk Manager...");
    
    pthread_mutex_lock(&dm_to_fini->lock);
    if (dm_to_fini->managed_disks) {
        xsan_list_destroy(dm_to_fini->managed_disks);
        dm_to_fini->managed_disks = NULL;
    }
    if (dm_to_fini->managed_disk_groups) {
        xsan_list_destroy(dm_to_fini->managed_disk_groups);
        dm_to_fini->managed_disk_groups = NULL;
    }
    pthread_mutex_unlock(&dm_to_fini->lock);
    pthread_mutex_destroy(&dm_to_fini->lock);

    if (dm_to_fini->md_store) {
        xsan_metadata_store_close(dm_to_fini->md_store);
        dm_to_fini->md_store = NULL;
    }

    XSAN_FREE(dm_to_fini);
    if (dm_ptr) *dm_ptr = NULL;
    if (dm_to_fini == g_xsan_disk_manager_instance) g_xsan_disk_manager_instance = NULL;
    
    XSAN_LOG_INFO("XSAN Disk Manager finalized.");
}

// Stub implementations for missing functions
xsan_disk_group_t *xsan_disk_manager_find_disk_group_by_id(xsan_disk_manager_t *dm, xsan_group_id_t group_id) {
    if (!dm || spdk_uuid_is_null((struct spdk_uuid*)&group_id.data[0])) {
        return NULL;
    }
    
    pthread_mutex_lock(&dm->lock);
    xsan_list_iter_t iter;
    xsan_list_iter_init(dm->managed_disk_groups, &iter);
    
    xsan_disk_group_t *group;
    while (xsan_list_iter_next(&iter, (void**)&group)) {
        if (spdk_uuid_compare((struct spdk_uuid*)&group->id.data[0], (struct spdk_uuid*)&group_id.data[0]) == 0) {
            pthread_mutex_unlock(&dm->lock);
            return group;
        }
    }
    
    pthread_mutex_unlock(&dm->lock);
    return NULL;
}

xsan_disk_t *xsan_disk_manager_find_disk_by_id(xsan_disk_manager_t *dm, xsan_disk_id_t disk_id) {
    if (!dm || spdk_uuid_is_null((struct spdk_uuid*)&disk_id.data[0])) {
        return NULL;
    }
    
    pthread_mutex_lock(&dm->lock);
    xsan_list_iter_t iter;
    xsan_list_iter_init(dm->managed_disks, &iter);
    
    xsan_disk_t *disk;
    while (xsan_list_iter_next(&iter, (void**)&disk)) {
        if (spdk_uuid_compare((struct spdk_uuid*)&disk->id.data[0], (struct spdk_uuid*)&disk_id.data[0]) == 0) {
            pthread_mutex_unlock(&dm->lock);
            return disk;
        }
    }
    
    pthread_mutex_unlock(&dm->lock);
    return NULL;
}

// 移除这些函数，因为它们与头文件声明不匹配
// 这些函数应该在头文件中声明为不同的签名

// 添加缺失的函数实现
xsan_error_t xsan_disk_manager_scan_and_register_bdevs(xsan_disk_manager_t *dm) {
    if (!dm || !dm->initialized) {
        return XSAN_ERROR_INVALID_PARAM;
    }
    if (spdk_get_thread() == NULL) {
        XSAN_LOG_ERROR("xsan_disk_manager_scan_and_register_bdevs must be called from an SPDK thread.");
        return XSAN_ERROR_THREAD_CONTEXT;
    }

    XSAN_LOG_INFO("Scanning and registering SPDK bdevs...");
    
    xsan_bdev_info_t *bdev_list = NULL;
    int bdev_count = 0;
    xsan_error_t err = xsan_bdev_list_get_all(&bdev_list, &bdev_count);
    if (err != XSAN_OK) {
        XSAN_LOG_ERROR("Failed to get SPDK bdev list: %s", xsan_error_to_string(err));
        return err;
    }

    pthread_mutex_lock(&dm->lock);
    
    // 清理现有磁盘列表（如果需要重新扫描）
    xsan_list_clear(dm->managed_disks);
    
    for (int i = 0; i < bdev_count; i++) {
        xsan_disk_t *disk = (xsan_disk_t *)XSAN_MALLOC(sizeof(xsan_disk_t));
        if (!disk) {
            XSAN_LOG_ERROR("Failed to allocate memory for disk %s", bdev_list[i].name);
            continue;
        }
        
        memset(disk, 0, sizeof(xsan_disk_t));
        
        // 生成UUID
        spdk_uuid_generate((struct spdk_uuid*)&disk->id.data[0]);
        
        // 复制基本信息
        xsan_strcpy_safe(disk->name, bdev_list[i].name, XSAN_MAX_NAME_LEN);
        xsan_strcpy_safe(disk->bdev_name, bdev_list[i].name, XSAN_MAX_NAME_LEN);
        disk->size_bytes = bdev_list[i].capacity_bytes;
        disk->block_size = bdev_list[i].block_size;
        disk->is_ssd = !bdev_list[i].is_rotational;
        disk->state = 1; // 假设为正常状态
        
        // 添加到列表
        if (xsan_list_append(dm->managed_disks, disk) != XSAN_OK) {
            XSAN_LOG_ERROR("Failed to add disk %s to managed list", disk->name);
            XSAN_FREE(disk);
        } else {
            char disk_id_str[SPDK_UUID_STRING_LEN];
            spdk_uuid_fmt_lower(disk_id_str, sizeof(disk_id_str), (struct spdk_uuid*)&disk->id.data[0]);
            XSAN_LOG_INFO("Registered disk: %s (ID: %s, Size: %lu bytes)", 
                         disk->name, disk_id_str, disk->size_bytes);
        }
    }
    
    pthread_mutex_unlock(&dm->lock);
    
    xsan_bdev_list_free(bdev_list, bdev_count);
    XSAN_LOG_INFO("SPDK bdev scan complete. Registered %d disks.", bdev_count);
    return XSAN_OK;
}

xsan_error_t xsan_disk_manager_disk_group_create(xsan_disk_manager_t *dm,
                                                 const char *group_name,
                                                 xsan_disk_group_type_t group_type,
                                                 const char *bdev_names_list[],
                                                 int num_bdevs,
                                                 xsan_group_id_t *group_id_out) {
    if (!dm || !dm->initialized || !group_name || !bdev_names_list || num_bdevs <= 0) {
        return XSAN_ERROR_INVALID_PARAM;
    }
    if (spdk_get_thread() == NULL) {
        XSAN_LOG_ERROR("xsan_disk_manager_disk_group_create must be called from an SPDK thread.");
        return XSAN_ERROR_THREAD_CONTEXT;
    }

    pthread_mutex_lock(&dm->lock);
    
    // 检查组名是否已存在
    xsan_list_iter_t iter;
    xsan_list_iter_init(dm->managed_disk_groups, &iter);
    xsan_disk_group_t *existing_group;
    while (xsan_list_iter_next(&iter, (void**)&existing_group)) {
        if (strcmp(existing_group->name, group_name) == 0) {
            pthread_mutex_unlock(&dm->lock);
            return XSAN_ERROR_ALREADY_EXISTS;
        }
    }
    
    // 创建新磁盘组
    xsan_disk_group_t *group = (xsan_disk_group_t *)XSAN_MALLOC(sizeof(xsan_disk_group_t));
    if (!group) {
        pthread_mutex_unlock(&dm->lock);
        return XSAN_ERROR_OUT_OF_MEMORY;
    }
    
    memset(group, 0, sizeof(xsan_disk_group_t));
    
    // 生成组UUID
    spdk_uuid_generate((struct spdk_uuid*)&group->id.data[0]);
    
    // 设置基本信息
    xsan_strcpy_safe(group->name, group_name, XSAN_MAX_NAME_LEN);
    group->type = group_type;
    group->disk_count = 0;
    group->total_capacity = 0;
    group->available_capacity = 0;
    group->state = 1; // 正常状态
    
    // 查找并添加磁盘
    for (int i = 0; i < num_bdevs && i < XSAN_MAX_REPLICAS; i++) {
        xsan_disk_t *disk = xsan_disk_manager_find_disk_by_bdev_name(dm, bdev_names_list[i]);
        if (disk) {
            group->disk_ids[group->disk_count] = disk->id;
            group->disk_count++;
            group->total_capacity += disk->size_bytes;
            group->available_capacity += disk->size_bytes;
        } else {
            XSAN_LOG_WARN("Bdev %s not found, skipping", bdev_names_list[i]);
        }
    }
    
    if (group->disk_count == 0) {
        XSAN_FREE(group);
        pthread_mutex_unlock(&dm->lock);
        return XSAN_ERROR_NOT_FOUND;
    }
    
    // 添加到列表
    if (xsan_list_append(dm->managed_disk_groups, group) != XSAN_OK) {
        XSAN_FREE(group);
        pthread_mutex_unlock(&dm->lock);
        return XSAN_ERROR_OUT_OF_MEMORY;
    }
    
    // 返回组ID
    if (group_id_out) {
        *group_id_out = group->id;
    }
    
    pthread_mutex_unlock(&dm->lock);
    
    char group_id_str[SPDK_UUID_STRING_LEN];
    spdk_uuid_fmt_lower(group_id_str, sizeof(group_id_str), (struct spdk_uuid*)&group->id.data[0]);
    XSAN_LOG_INFO("Created disk group: %s (ID: %s, Type: %d, Disks: %d, Capacity: %lu bytes)", 
                 group->name, group_id_str, group->type, group->disk_count, group->total_capacity);
    
    return XSAN_OK;
}

xsan_error_t xsan_disk_manager_disk_group_delete(xsan_disk_manager_t *dm, xsan_group_id_t group_id) {
    if (!dm || !dm->initialized) {
        return XSAN_ERROR_INVALID_PARAM;
    }
    if (spdk_get_thread() == NULL) {
        XSAN_LOG_ERROR("xsan_disk_manager_disk_group_delete must be called from an SPDK thread.");
        return XSAN_ERROR_THREAD_CONTEXT;
    }

    pthread_mutex_lock(&dm->lock);
    
    xsan_list_iter_t iter;
    xsan_list_iter_init(dm->managed_disk_groups, &iter);
    
    xsan_disk_group_t *group;
    while (xsan_list_iter_next(&iter, (void**)&group)) {
        if (spdk_uuid_compare((struct spdk_uuid*)&group->id.data[0], (struct spdk_uuid*)&group_id.data[0]) == 0) {
            // 找到要删除的组
            char group_id_str[SPDK_UUID_STRING_LEN];
            spdk_uuid_fmt_lower(group_id_str, sizeof(group_id_str), (struct spdk_uuid*)&group->id.data[0]);
            XSAN_LOG_INFO("Deleting disk group: %s (ID: %s)", group->name, group_id_str);
            
            // 从列表中移除
            xsan_list_iter_remove(&iter);
            pthread_mutex_unlock(&dm->lock);
            return XSAN_OK;
        }
    }
    
    pthread_mutex_unlock(&dm->lock);
    return XSAN_ERROR_NOT_FOUND;
}

xsan_disk_t *xsan_disk_manager_find_disk_by_bdev_name(xsan_disk_manager_t *dm, const char *bdev_name) {
    if (!dm || !bdev_name) {
        return NULL;
    }
    
    pthread_mutex_lock(&dm->lock);
    xsan_list_iter_t iter;
    xsan_list_iter_init(dm->managed_disks, &iter);
    
    xsan_disk_t *disk;
    while (xsan_list_iter_next(&iter, (void**)&disk)) {
        if (strcmp(disk->bdev_name, bdev_name) == 0) {
            pthread_mutex_unlock(&dm->lock);
            return disk;
        }
    }
    
    pthread_mutex_unlock(&dm->lock);
    return NULL;
}

xsan_error_t xsan_disk_group_allocate_extents(xsan_disk_manager_t *dm,
                                              xsan_group_id_t group_id,
                                              uint64_t num_blocks,
                                              uint32_t block_size,
                                              xsan_volume_extent_mapping_t **extents_out,
                                              uint32_t *num_extents_out) {
    if (!dm || !extents_out || !num_extents_out || num_blocks == 0 || block_size == 0) {
        return XSAN_ERROR_INVALID_PARAM;
    }
    
    // Stub implementation - just return success with empty extents
    *extents_out = NULL;
    *num_extents_out = 0;
    return XSAN_OK;
}

xsan_error_t xsan_disk_group_free_extents(xsan_disk_manager_t *dm,
                                          xsan_group_id_t group_id,
                                          const xsan_volume_extent_mapping_t *extents,
                                          uint32_t num_extents) {
    if (!dm || !extents || num_extents == 0) {
        return XSAN_ERROR_INVALID_PARAM;
    }
    
    // Stub implementation - just return success
    return XSAN_OK;
}
